import argparse
import os
import pickle
import warnings
import xml.etree.ElementTree as ET

import cv2
import matplotlib.pyplot as plt
import numpy as np
import scipy.io as sio
import spectral.io.envi as envi
from linformer import Linformer
from numpy import flip
from sklearn import model_selection, svm
from sklearn.neighbors import KNeighborsClassifier

from utils import (DataPreProcess, DataResult, get_imggnd, listen_plot, myplot,
                   plot, setpath, splitdata)

# from vit_pytorch.efficient import ViT

def mysvm(x_train, y_train, x_test):
    SVM_GRID_PARAMS = [{'kernel': ['rbf'], 'gamma': [1e-1, 1e-2, 1e-3],
                                       'C': [1, 10, 100, 1000]},
                #    {'kernel': ['linear'], 'C': [0.1, 1, 10, 100, 1000]},
                #    {'kernel': ['poly'], 'degree': [3], 'gamma': [1e-1, 1e-2, 1e-3]}
                ]
    class_weight = 'balanced'
   
    try:
        clf = svm.SVC(class_weight=class_weight)
        clf = model_selection.GridSearchCV(clf, SVM_GRID_PARAMS, verbose=1, n_jobs=4)
        clf.fit(x_train, y_train)
    except ValueError:
        clf = svm.SVC(class_weight=class_weight)
        clf.fit(x_train, y_train)
    # print("SVM best parameters : {}".format(clf.best_params_))
    prediction = clf.predict(x_test)
    return prediction

def myknn(x_train, y_train, x_test):
    kclf = KNeighborsClassifier(n_neighbors=1)
    kclf.fit(x_train, y_train)
    prediction = kclf.predict(x_test)   
    return prediction 





if __name__ == '__main__':
    print('='*30)
    warnings.filterwarnings("ignore")
    parser = argparse.ArgumentParser()
    warnings.filterwarnings("ignore")
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str,  default='./pathology/data/032370b-20x-roi2', metavar='D',
                        help='the dataset path you load')
    parser.add_argument('--trial_number', type=int,  default=1, metavar='T',
                        help='the time you do this trial')
    parser.add_argument('--train_number', type=int,  default=0.01, metavar='NTr', 
                        help='number of training set')
    parser.add_argument('--valid_number', type=int,  default=0.01, metavar='NVa',
                        help='number of valid set')
    parser.add_argument('--test_number', type=int,  default=0.98, metavar='NTe',
                        help='number of test set')
    parser.add_argument('--patchsize', type=int,  default=9, metavar='P',
                        help='patchsize of data') 
    parser.add_argument('--modelname', type=str,  default='SVM', metavar='P', help='which model to choose') 
    parser.add_argument('--gpu_ids', type=int,  default=3, metavar='G',
                        help='which gpu to use')
    

    args = parser.parse_args()
   
    dataset = args.dataset
    NTr = args.train_number
    trialnumber = args.trial_number
    NTe = args.test_number
    NVa = args.valid_number
    patchsize = args.patchsize
    modelname = args.modelname

   


    resultpath, imagepath, datapath = setpath(dataset, trialnumber , NTr, NVa, NTe, modelname)
    
    IMAGE, GND = get_imggnd(dataset)

    
    spliteddata = splitdata(IMAGE, GND, datapath , trainnum=NTr, validnum=NVa,testnum=NTe)

    processeddata = DataPreProcess(IMAGE, patchsize, datapath, 1).processeddata
    x_train = processeddata['train'].patch.squeeze()
    x_valid = processeddata['valid'].patch.squeeze()
    x_test = processeddata['test'].patch.squeeze()
    y_train = processeddata['train'].gt
    y_valid = processeddata['valid'].gt
    y_test = processeddata['test'].gt
    x_train = np.vstack((x_train, x_valid))
    x_train = x_train[:, 4, 4, :].squeeze()
    y_train = np.hstack((y_train, y_valid))
    x_test = x_test[:, 4, 4, :].squeeze()
    print("Running a grid search SVM")
# Grid search SVM (linear and RBF)
    # knn_prediction = myknn(x_train, y_train, x_test)
    # print(sum(knn_prediction==y_test)/y_test.shape[0])
    svm_prediction = mysvm(x_train, y_train, x_test)
    D = DataResult()
    D.y_score = svm_prediction
    D.y_true = y_test
    D.get_metric()
    with open(resultpath + 'result.pkl', 'wb') as f:
        pickle.dump(D, f, pickle.HIGHEST_PROTOCOL)
    imgPos = np.array(list(processeddata['train'].pos) + list(processeddata['valid'].pos) + list(processeddata['test'].pos))
    imgGt = np.array(list(processeddata['train'].gt) + list(processeddata['valid'].gt) + list(svm_prediction))
    plot(imgPos, imgGt, IMAGE.shape, imagepath + 'svmprediction')
    # imgGt = np.array(list(processeddata['train'].gt) + list(processeddata['valid'].gt) + list(knn_prediction))
    # plot(imgPos, imgGt, IMAGE.shape, imagepath + 'knnprediction')

  
